#!/bin/env python

from plot_calc_functions import plot_out
import matplotlib
matplotlib.use('Agg')
matplotlib.rc('xtick',labelsize='x-large')
matplotlib.rc('ytick',labelsize='x-large')
# matplotlib.rc('axes.formatter',limits='-2,3')

from pylab import *
import sys, os
import pylab as p
import numpy as np

## This takes in a text file of run directories, particle ids, and legendlabels, makes multiple plots ##

############# Change here #############

fMbh = {'plotcalc':'None','colx':1,'coly':2,'xlog':1,'ylog':1,'title':'Black hole mass','saveas':'fMbh','xlim':(1,12),'ylim':(1e5,1e12),
	'xlabel':r'$(1+z)$','ylabel':r'$\mathrm{M}_{\mathrm{BH}}$ [$\mathrm{M}_{\odot}$]', 'legendposition':'best'}
fcube_mass = {}
fx = {'plotcalc':'None','colx':1,'coly':4,'xlog':1,'ylog':0,'title':'Black hole position','saveas':'fx','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$x$ [code]', 'legendposition':'best'}
fy = {'plotcalc':'None','colx':5,'coly':6,'xlog':0,'ylog':0,'title':'Black hole position','saveas':'fy','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$y$ [code]', 'legendposition':'best'}
fz = {'plotcalc':'None','colx':6,'coly':4,'xlog':0,'ylog':0,'title':'Black hole position','saveas':'fz','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$z$ [code]', 'legendposition':'best'}
fMbhl = {'plotcalc':'None','colx':1,'coly':10,'xlog':1,'ylog':1,'title':'Bondi Rate','saveas':'fMbhl','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$\dot{\mathrm{M}}_{\mathrm{Bondi}}$ [$\mathrm{M}_{\odot}/\mathrm{Myr}$]', 'legendposition':'best'}
fMacc = {'plotcalc':'None','colx':1,'coly':11,'xlog':1,'ylog':1,'title':'BH Accretion Rate','saveas':'fMacc','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$\dot{\mathrm{M}}_{\mathrm{BH}}$ [$\mathrm{M}_{\odot}/\mathrm{Myr}$]', 'legendposition':'best'}
fMedd = {'plotcalc':'None','colx':1,'coly':12,'xlog':1,'ylog':1,'title':'Eddington Rate','saveas':'fMedd','xlim':(1,12),'ylim':'None',
	'xlabel':r'$(1+z)$','ylabel':r'$\dot{\mathrm{M}}_{\mathrm{Edd}}$ [$\mathrm{M}_{\odot}/\mathrm{Myr}$]', 'legendposition':'best'}
fT = {'plotcalc':'None','colx':1,'coly':13,'xlog':1,'ylog':1,'title':'Temperature','saveas':'fT','xlim':(1,12),'ylim':(1,1e10),
      'xlabel':r'$(1+z)$','ylabel':r'$\mathrm{T}_{\mathrm{sink}}$ [K]','legendposition':'best'}
fcs = {'plotcalc':'None','colx':1,'coly':15,'xlog':1,'ylog':1,'title':'Sound Speed','saveas':'fcs','xlim':(1,12),'ylim':(1,1e6),
      'xlabel':r'$(1+z)$','ylabel':r'$\mathrm{c}_{\mathrm{s,local}}$ [km/s]','legendposition':'best'}
fnH = {'plotcalc':'None','colx':1,'coly':17,'xlog':1,'ylog':1,'title':'Hydrogen Number Density','saveas':'fnH','xlim':(1,12),'ylim':(1e-3,1e5),
      'xlabel':r'$(1+z)$','ylabel':r'$\mathrm{n}_{\mathrm{H}}$ [$\mathrm{cm}^{-3}$]','legendposition':'best'}
fdv = {}

plot_items = [ fMbhl, fT, fcs, fx, fy, fz, fnH, fMacc # , fMbhl, fx, fy, fz, fcs, fT, fMacc, fnH
	       ] #fMbh, fx, fy, fz, fMbhl, fMacc, fcs, fT, fnH ]

#######################################

class MyLogFormatter(LogFormatter) :

	def __call__(self,x,pos=None) :
		if abs(log10(x)) <2 :
			return "$%g$" % (x,)
		else :
			return "$10^{%g}$" % (log10(x),)

if len(sys.argv) != 2 :
	print 'Usage: plot.py listofdirectories.txt[rundir, BHid, legendlabel,markercolor/type]'
	sys.exit(1)

rundirlist = sys.argv[1]
png_folder = rundirlist.split('.txt')[0].split('/')[-1]
print png_folder
if not os.path.exists( png_folder ) :
	print 'Making '+png_folder
	os.mkdir(png_folder)
rundirlines = [ line for line in open(rundirlist,'r').read().splitlines() ]
legend_labels = rundirlines[len(rundirlines)/2:3*len(rundirlines)/4]
rundirs = dict(zip(legend_labels,rundirlines[:len(rundirlines)/4]))
bh_ids = dict(zip(legend_labels,rundirlines[len(rundirlines)/4:len(rundirlines)/2]))
markers = dict(zip(legend_labels,rundirlines[3*len(rundirlines)/4:]))

### PLOT EACH PLOT IN LIST OF PLOT ITEMS ###
		   
for plot_item in plot_items :
	# Create a 5x5 figure
	f = p.figure( figsize = (5,5) )
	axes([0.20,0.15,0.75,0.75])
	plots = []
	# Loop over each data set
	for legend_label in legend_labels :
		xdata, ydata = plot_out( rundirs[legend_label],
					 bh_ids[legend_label],
					 plot_item )  
		plots.append(f.gca().plot( xdata, ydata, markers[legend_label], markersize=0.5 ))

	# Specify plot settings
	f.gca().set_xlabel(plot_item['xlabel'],fontsize='xx-large')
	f.gca().set_ylabel(plot_item['ylabel'],fontsize='xx-large')
	if (plot_item['xlog']) :
		f.gca().set_xscale('log')
		f.gca().xaxis.set_major_formatter(MyLogFormatter())
	if (plot_item['ylog']) :
		f.gca().set_yscale('log')
		f.gca().yaxis.set_major_formatter(MyLogFormatter())
	if plot_item['xlim'] != 'None' :
		f.gca().set_xlim(plot_item['xlim'])
	if plot_item['ylim'] != 'None' :
		f.gca().set_ylim(plot_item['ylim'])

	f.gca().legend(plots,legend_labels,plot_item['legendposition'],markerscale=4.)
	leg = gca().get_legend()
	legtext = leg.get_texts()
	leglines = leg.get_lines()
	setp(legtext,fontsize='large')
	setp(leglines,linewidth=4.)
	leg.draw_frame(False)

	f.gca().set_title(plot_item['title'],fontsize='xx-large')

	# Save the figure
	f.savefig( png_folder+'/'+plot_item['saveas'], dpi=100 )








